import torch
import torch.nn as nn

class InConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(InConv, self).__init__()
        self.conv = DoubleConv(in_ch, out_ch)

    def forward(self, x):
        x = self.conv(x)
        return x

class Downsample(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(Downsample, self).__init__()
        self.mpconv = nn.Sequential(
            nn.MaxPool3d(2, 2),
            DoubleConv(in_ch, out_ch)
        )

    def forward(self, x):
        x = self.mpconv(x)
        return x

class OutConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(OutConv, self).__init__()
        self.conv = nn.Conv3d(in_ch, out_ch, 1)
        # self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.conv(x)
        # x = self.sigmoid(x)
        return x

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv3d(in_ch, out_ch, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm3d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv3d(out_ch, out_ch, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm3d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x

class Upsample(nn.Module):
    def __init__(self, in_ch, skip_ch,out_ch):
        super(Upsample, self).__init__()
        self.up = nn.ConvTranspose3d(in_ch, in_ch, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_ch+skip_ch, out_ch)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x
class ChannelAttention(nn.Module):
    """Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet."""

    def __init__(self, channels: int) -> None:
        """Initializes the class and sets the basic configurations and instance variables required."""
        super().__init__()
        self.pool = nn.AdaptiveAvgPool3d(1)
        self.fc = nn.Conv3d(channels, channels, 1, 1, 0, bias=True)
        self.act = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Applies forward pass using activation on convolutions of the input, optionally using batch normalization."""
        return x * self.act(self.fc(self.pool(x)))


class TransformerLayer(nn.Module):
    # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
    def __init__(self, c, num_heads):
        """
        Initializes a transformer layer, sans LayerNorm for performance, with multihead attention and linear layers.

        See  as described in https://arxiv.org/abs/2010.11929.
        """
        super().__init__()
        self.q = nn.Linear(c, c, bias=False)
        self.k = nn.Linear(c, c, bias=False)
        self.v = nn.Linear(c, c, bias=False)
        self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
        self.fc1 = nn.Linear(c, c, bias=False)
        self.fc2 = nn.Linear(c, c, bias=False)

    def forward(self, x):
        """Performs forward pass using MultiheadAttention and two linear transformations with residual connections."""
        b,c,t,h,w = x.shape
        p = x.flatten(2).permute(2,0,1)
        p = self.ma(self.q(p), self.k(p), self.v(p))[0] + p
        p = self.fc2(self.fc1(p)) + p
        return p.permute(1,2,0).reshape(b,c,t,h,w)

class CMUNet(nn.Module):
    def __init__(self, in_channels, num_classes):
        super(CMUNet, self).__init__()
        features = [32,64,128,256]

        self.inc = InConv(in_channels, features[0])
        self.down1 = Downsample(features[0], features[1])
        self.down2 = Downsample(features[1], features[2])
        self.down3 = Downsample(features[2], features[3])
        self.down4 = Downsample(features[3], features[3])

        self.up1 = Upsample(features[3], features[3], features[2])
        self.up2 = Upsample(features[2], features[2], features[1])
        self.up3 = Upsample(features[1], features[1], features[0])
        self.up4 = Upsample(features[0], features[0], features[0])

        self.ca = ChannelAttention(features[3])#
        #self.ca1 = ChannelAttention(features[1])
        self.sa = TransformerLayer(features[3],2)#
        self.outc = OutConv(features[0], num_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x5 = self.ca(x5)
        x5 = self.sa(x5)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        #x = self.ca1(x)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        x = self.outc(x)

        return x


if __name__ == '__main__':
    from fvcore.nn import FlopCountAnalysis, parameter_count_table

    # x = torch.rand(2, 4, 32, 32, 16)
    x = torch.rand(1, 4, 160, 160, 128)
    print('x.shape:', x.shape)
    net = CMUNet(in_channels=4, num_classes=4)
    y = net(x)
    flops = FlopCountAnalysis(net, x)
    print('flops:', flops.total())
    print('params:', parameter_count_table(net))
    print("params: ", sum(p.numel() for p in net.parameters()))
    print(y.shape)